//! 沙箱生命周期管理模块
//! 
//! 提供沙箱从创建到销毁的完整生命周期管理

use crate::{Result, Error, ErrorContext};
use crate::types::{SandboxState, SandboxMetadata, Language, SandboxSpec, SandboxId};
use crate::firecracker::vm::FirecrackerVm;
use crate::logging::{LoggingDecorator, LogContext};
use crate::retry::{RetryExecutor, RetryPolicy};
use std::time::{Duration, Instant};
use std::sync::Arc;
use dashmap::DashMap;
use tokio::time::{sleep, timeout};
use tokio::sync::Mutex;
use tracing::{warn, error, debug};

/// 生命周期管理器
#[derive(Clone)]
pub struct LifecycleManager {
    /// 活跃沙箱映射
    active_sandboxes: Arc<DashMap<SandboxId, Arc<Mutex<FirecrackerVm>>>>,
    /// 重试执行器
    retry_executor: RetryExecutor,
    /// 配置
    config: Arc<crate::Config>,
}

impl LoggingDecorator for LifecycleManager {
    fn component_name(&self) -> &'static str {
        "lifecycle_manager"
    }
}

impl LifecycleManager {
    /// 创建生命周期管理器
    pub fn new(config: Arc<crate::Config>) -> Self {
        Self {
            active_sandboxes: Arc::new(DashMap::new()),
            retry_executor: RetryExecutor::new(RetryPolicy::standard()),
            config,
        }
    }
    
    /// 创建沙箱
    /// 
    /// 创建一个新的沙箱实例，包括VM初始化和资源分配
    pub async fn create_sandbox(&self, spec: SandboxSpec) -> Result<SandboxMetadata> {
        let start_time = Instant::now();
        let mut metadata = SandboxMetadata::new(spec.clone());
        
        self.log_info("创建沙箱开始", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .with_str("language", &spec.language.to_string())
                .with_number("memory_mb", spec.memory_mb)
                .with_number("cpu_cores", spec.cpu_cores)
                .to_json()
        ));
        
        // 使用重试机制创建VM
        let vm_result = self.retry_executor.execute(
            "create_firecracker_vm",
            {
                let id = metadata.id;
                let spec = spec.clone();
                move || {
                    let id = id;
                    let spec = spec.clone();
                    async move {
                        Self::create_vm_internal(id, spec).await
                    }
                }
            }
        ).await;
        
        match vm_result {
            Ok(vm) => {
                // 存储VM实例
                self.active_sandboxes.insert(metadata.id, Arc::new(Mutex::new(vm)));
                
                // 更新元数据状态
                metadata.state = SandboxState::Stopped;
                
                let duration = start_time.elapsed();
                self.log_info("沙箱创建成功", Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_number("duration_ms", duration.as_millis() as i64)
                        .to_json()
                ));
                
                Ok(metadata)
            }
            Err(e) => {
                // 标记为错误状态
                metadata.state = SandboxState::Error(e.to_string());
                
                self.log_error(&e, Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_str("operation", "create_sandbox")
                        .to_json()
                ));
                
                Err(e).context("沙箱创建失败")
            }
        }
    }
    
    /// 启动沙箱
    /// 
    /// 启动已创建的沙箱，使其进入可用状态
    pub async fn start_sandbox(&self, metadata: &mut SandboxMetadata) -> Result<()> {
        let start_time = Instant::now();
        
        // 检查当前状态
        if !matches!(metadata.state, SandboxState::Stopped) {
            return Err(Error::sandbox_start(format!(
                "沙箱状态不正确，当前状态: {:?}", 
                metadata.state
            )));
        }
        
        // 更新状态为启动中
        metadata.state = SandboxState::Starting;
        
        self.log_info("启动沙箱", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .to_json()
        ));
        
        // 获取VM实例
        let vm = self.active_sandboxes.get(&metadata.id)
            .ok_or_else(|| Error::sandbox_not_found(metadata.id.to_string()))?;
        
        // 使用重试机制启动VM
        let start_result = self.retry_executor.execute(
            "start_firecracker_vm",
            {
                let vm = vm.clone();
                move || {
                    let vm = vm.clone();
                    async move {
                        let mut vm_guard = vm.lock().await;
                        vm_guard.start().await
                    }
                }
            }
        ).await;
        
        match start_result {
            Ok(()) => {
                // 等待VM完全启动
                self.wait_for_vm_ready(vm.clone(), Duration::from_secs(30)).await?;
                
                // 更新元数据
                metadata.mark_started();
                
                let duration = start_time.elapsed();
                self.log_info("沙箱启动成功", Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_number("duration_ms", duration.as_millis() as i64)
                        .to_json()
                ));
                
                Ok(())
            }
            Err(e) => {
                // 标记为错误状态
                metadata.state = SandboxState::Error(e.to_string());
                
                self.log_error(&e, Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_str("operation", "start_sandbox")
                        .to_json()
                ));
                
                Err(e).context("沙箱启动失败")
            }
        }
    }
    
    /// 停止沙箱
    /// 
    /// 优雅地停止沙箱，保持数据完整性
    pub async fn stop_sandbox(&self, metadata: &mut SandboxMetadata) -> Result<()> {
        let start_time = Instant::now();
        
        // 检查当前状态
        if matches!(metadata.state, SandboxState::Stopped | SandboxState::Error(_)) {
            return Ok(()); // 已经停止或错误状态，无需操作
        }
        
        // 如果正在执行任务，等待任务完成或超时
        if matches!(metadata.state, SandboxState::Busy) {
            self.log_info("等待沙箱任务完成", Some(
                LogContext::new()
                    .with_str("sandbox_id", &metadata.id.to_string())
                    .to_json()
            ));
            
            // 等待最多30秒让任务完成
            let wait_timeout = Duration::from_secs(30);
            let wait_start = Instant::now();
            
            while matches!(metadata.state, SandboxState::Busy) && wait_start.elapsed() < wait_timeout {
                sleep(Duration::from_millis(100)).await;
            }
        }
        
        // 更新状态为停止中
        metadata.state = SandboxState::Stopping;
        
        self.log_info("停止沙箱", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .to_json()
        ));
        
        // 获取VM实例
        if let Some(vm) = self.active_sandboxes.get(&metadata.id) {
            // 使用重试机制停止VM
            let stop_result = self.retry_executor.execute(
                "stop_firecracker_vm",
                {
                    let vm = vm.clone();
                    move || {
                        let vm = vm.clone();
                        async move {
                            let mut vm_guard = vm.lock().await;
                            vm_guard.stop().await
                        }
                    }
                }
            ).await;
            
            if let Err(e) = stop_result {
                self.log_error(&e, Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_str("operation", "stop_sandbox")
                        .to_json()
                ));
                
                // 即使停止失败，也继续执行清理流程
                warn!("VM停止失败，但继续执行清理: {}", e);
            }
        }
        
        // 更新元数据
        metadata.state = SandboxState::Stopped;
        
        let duration = start_time.elapsed();
        self.log_info("沙箱停止成功", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .with_number("duration_ms", duration.as_millis() as i64)
                .to_json()
        ));
        
        Ok(())
    }
    
    /// 销毁沙箱
    /// 
    /// 完全销毁沙箱，释放所有资源
    pub async fn destroy_sandbox(&self, metadata: &mut SandboxMetadata) -> Result<()> {
        let start_time = Instant::now();
        
        self.log_info("销毁沙箱", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .to_json()
        ));
        
        // 首先停止沙箱（如果还在运行）
        if !matches!(metadata.state, SandboxState::Stopped | SandboxState::Error(_)) {
            if let Err(e) = self.stop_sandbox(metadata).await {
                warn!("停止沙箱时出错，但继续销毁流程: {}", e);
            }
        }
        
        // 从活跃沙箱映射中移除
        if let Some((_, vm)) = self.active_sandboxes.remove(&metadata.id) {
            // 停止并销毁VM资源
            let stop_result = vm.lock().await.stop().await;
            
            if let Err(e) = stop_result {
                self.log_error(&e, Some(
                    LogContext::new()
                        .with_str("sandbox_id", &metadata.id.to_string())
                        .with_str("operation", "destroy_vm")
                        .to_json()
                ));
                
                warn!("VM资源销毁失败: {}", e);
            }
        }
        
        let duration = start_time.elapsed();
        self.log_info("沙箱销毁完成", Some(
            LogContext::new()
                .with_str("sandbox_id", &metadata.id.to_string())
                .with_number("duration_ms", duration.as_millis() as i64)
                .with_number("total_usage_count", metadata.usage_count as i64)
                .to_json()
        ));
        
        Ok(())
    }
    
    /// 获取沙箱健康状态
    /// 
    /// 检查沙箱是否健康运行
    pub async fn get_sandbox_health(&self, sandbox_id: SandboxId) -> Result<bool> {
        if let Some(vm) = self.active_sandboxes.get(&sandbox_id) {
            let vm_guard = vm.lock().await;
            Ok(vm_guard.is_running())
        } else {
            Ok(false)
        }
    }
    
    /// 执行沙箱健康检查
    /// 
    /// 对所有活跃沙箱执行健康检查
    pub async fn perform_health_check(&self) -> Result<()> {
        let total_sandboxes = self.active_sandboxes.len();
        let mut healthy_count = 0;
        let mut unhealthy_sandboxes = Vec::new();
        
        for entry in self.active_sandboxes.iter() {
            let sandbox_id = *entry.key();
            match self.get_sandbox_health(sandbox_id).await {
                Ok(true) => healthy_count += 1,
                Ok(false) => unhealthy_sandboxes.push(sandbox_id),
                Err(e) => {
                    warn!("健康检查失败 sandbox_id={}: {}", sandbox_id, e);
                    unhealthy_sandboxes.push(sandbox_id);
                }
            }
        }
        
        self.log_info("健康检查完成", Some(
            LogContext::new()
                .with_number("total_sandboxes", total_sandboxes as i64)
                .with_number("healthy_count", healthy_count as i64)
                .with_number("unhealthy_count", unhealthy_sandboxes.len() as i64)
                .to_json()
        ));
        
        // 清理不健康的沙箱
        for unhealthy_id in unhealthy_sandboxes {
            if let Some((_, vm)) = self.active_sandboxes.remove(&unhealthy_id) {
                warn!("清理不健康的沙箱: {}", unhealthy_id);
                
                // 在单独的作用域中处理VM停止，避免跨越await边界持有锁
                let stop_result = vm.lock().await.stop().await;
                
                if let Err(e) = stop_result {
                    error!("清理不健康沙箱失败: {}", e);
                }
            }
        }
        
        Ok(())
    }
    
    /// 获取活跃沙箱数量
    pub fn active_sandbox_count(&self) -> usize {
        self.active_sandboxes.len()
    }
    
    /// 获取特定语言的活跃沙箱数量
    pub async fn active_sandbox_count_by_language(&self, language: &Language) -> usize {
        let mut count = 0;
        
        for entry in self.active_sandboxes.iter() {
            let vm_guard = entry.value().lock().await;
            if &vm_guard.spec.language == language {
                count += 1;
            }
        }
        
        count
    }
    
    // 内部方法
    
    /// 创建VM的内部实现
    async fn create_vm_internal(id: SandboxId, spec: SandboxSpec) -> Result<FirecrackerVm> {
        // 创建VM实例
        let vm = FirecrackerVm::new(id, spec);
        
        Ok(vm)
    }
    
    /// 等待VM就绪
    async fn wait_for_vm_ready(&self, vm: Arc<Mutex<FirecrackerVm>>, timeout_duration: Duration) -> Result<()> {
        let timeout_result = timeout(timeout_duration, async {
            loop {
                let vm_guard = vm.lock().await;
                if vm_guard.is_running() {
                    debug!("VM已就绪");
                    break;
                }
                drop(vm_guard);
                sleep(Duration::from_millis(500)).await;
            }
        }).await;
        
        timeout_result.map_err(|_| Error::timeout("等待VM就绪超时"))
    }
}

impl Default for LifecycleManager {
    fn default() -> Self {
        Self::new(Arc::new(crate::Config::default()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{SandboxSpec, Language};
    
    #[tokio::test]
    async fn test_lifecycle_manager_creation() {
        let config = Arc::new(crate::Config::default());
        let manager = LifecycleManager::new(config);
        
        assert_eq!(manager.active_sandbox_count(), 0);
    }
    
    #[tokio::test]
    async fn test_sandbox_metadata_lifecycle() {
        let spec = SandboxSpec {
            language: Language::Python,
            ..Default::default()
        };
        
        let mut metadata = SandboxMetadata::new(spec);
        assert_eq!(metadata.state, SandboxState::Creating);
        
        // 模拟启动过程
        metadata.state = SandboxState::Starting;
        metadata.mark_started();
        assert_eq!(metadata.state, SandboxState::Running);
        assert!(metadata.started_at.is_some());
        
        // 模拟使用
        metadata.mark_busy();
        assert_eq!(metadata.state, SandboxState::Busy);
        assert_eq!(metadata.usage_count, 1);
        
        metadata.mark_available();
        assert_eq!(metadata.state, SandboxState::Running);
    }
    
    #[test]
    fn test_sandbox_metadata_timing() {
        let spec = SandboxSpec::default();
        let mut metadata = SandboxMetadata::new(spec);
        
        // 模拟启动
        metadata.mark_started();
        assert!(metadata.uptime().is_some());
        
        // 模拟使用
        metadata.mark_busy();
        metadata.mark_available();
        assert!(metadata.idle_time().is_some());
    }
}